#!/usr/bin/env python3
"""
简单ChromaDB功能测试
"""

import sys
import os

def test_basic_chromadb():
    """测试基本ChromaDB功能"""
    print("=== 基本ChromaDB功能测试 ===")
    
    try:
        # 检查ChromaDB是否可导入
        import chromadb
        print("✅ ChromaDB 可成功导入")
        print(f"   版本: {chromadb.__version__}")
        
        # 尝试创建内存客户端（避免网络连接问题）
        from chromadb.config import Settings
        
        client = chromadb.Client(Settings(
            chroma_db_impl="duckdb+parquet",
            persist_directory="./test_chroma_db"
        ))
        print("✅ ChromaDB内存客户端创建成功")
        
        # 创建测试集合
        collection = client.create_collection("simple_test")
        print("✅ 测试集合创建成功")
        
        # 添加文档
        collection.add(
            documents=["这是一个测试文档", "另一个测试文档"],
            metadatas=[{"type": "test"}, {"type": "test"}],
            ids=["doc1", "doc2"]
        )
        print("✅ 文档添加成功")
        
        # 查询
        results = collection.query(query_texts=["测试"], n_results=2)
        found_count = len(results['documents'][0]) if results['documents'] else 0
        print(f"✅ 查询成功，找到 {found_count} 个文档")
        
        # 清理
        client.delete_collection("simple_test")
        if os.path.exists("./test_chroma_db"):
            import shutil
            shutil.rmtree("./test_chroma_db")
        print("✅ 清理完成")
        
        return True
        
    except ImportError:
        print("❌ ChromaDB 未安装")
        return False
    except Exception as e:
        print(f"❌ ChromaDB 功能测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_system_detection():
    """测试系统对ChromaDB的检测"""
    print("\n=== 系统ChromaDB检测测试 ===")
    
    try:
        from utils.chroma_db import HAS_CHROMADB
        
        print(f"HAS_CHROMADB 检测结果: {HAS_CHROMADB}")
        
        if HAS_CHROMADB:
            print("✅ 系统正确检测到ChromaDB已安装")
            
            # 测试KnowledgeChromaDB初始化
            from utils.chroma_db import KnowledgeChromaDB
            
            kb = KnowledgeChromaDB("detection_test")
            print(f"✅ KnowledgeChromaDB初始化成功，模式: {kb.mode}")
            
            if kb.mode == "chromadb":
                print("✅ 系统正确使用ChromaDB模式")
            else:
                print("❌ 系统未使用ChromaDB模式")
                return False
                
            return True
        else:
            print("❌ 系统未检测到ChromaDB安装")
            return False
            
    except Exception as e:
        print(f"❌ 系统检测测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    print("🔍 开始简单ChromaDB测试")
    
    # 运行测试
    basic_ok = test_basic_chromadb()
    detection_ok = test_system_detection()
    
    # 输出结果
    print("\n" + "="*60)
    print("📋 简单测试结果:")
    print(f"   基本功能: {'✅ 成功' if basic_ok else '❌ 失败'}")
    print(f"   系统检测: {'✅ 成功' if detection_ok else '❌ 失败'}")
    
    if basic_ok and detection_ok:
        print("\n🎉 ChromaDB基本功能正常！")
        print("💡 网络连接问题不影响本地功能使用")
        sys.exit(0)
    else:
        print("\n❌ 测试失败")
        sys.exit(1)